package team.seekers.hula.system.ratelimit.limiter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;

import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
import java.util.stream.Stream;


/**
 * @author GH
 */
@Slf4j
public class RequestRateLimiter {

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    private AtomicBoolean initialized = new AtomicBoolean(false);

    private RedisScript<List<Long>> script;

    public RequestRateLimiter(RedisScript<List<Long>> script) {
        this.script = script;
        initialized.compareAndSet(false, true);
    }

    @SuppressWarnings("unchecked")
    public Boolean isAllowed(String routeId, String id, Map<String, Integer> config, int requested) {
        if (!this.initialized.get()) {
            throw new IllegalStateException("RedisRateLimiter is not initialized");
        }

        int replenishRate = Optional.ofNullable(config.get(routeId))
                .orElseGet(() -> config.getOrDefault("default", 10000));
        int capacity = Math.max(replenishRate, requested);
        try {
            Object[] scriptArgs = Stream.of(replenishRate, capacity, Instant.now().getEpochSecond(), requested)
                    .map(String::valueOf)
                    .toArray();
            return Optional.ofNullable(this.redisTemplate.execute(this.script, getKeys(id), scriptArgs))
                    .orElse(Arrays.asList(1L, -1L)).get(0) == 1L;
        } catch (Exception e) {
            log.error("通过redis判断用户是否被限流时出错", e);
        }
        return true;
    }

    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }
}
